import json, labelImg
import os
import os.path as osp
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np
import json
import collections, glob, uuid, pycocotools.mask, cv2


def write_xml(imgname,filepath,labeldicts):                     #参数imagename是图片名（无后缀）
    root = ET.Element('Annotation')                             #创建Annotation根节点
    ET.SubElement(root, 'filename').text = str(imgname)         #创建filename子节点（无后缀）
    sizes = ET.SubElement(root,'size')                          #创建size子节点            
    ET.SubElement(sizes, 'width').text = str(labeldicts[0]['img_w'])                #没带脑子直接写了原图片的尺寸......
    ET.SubElement(sizes, 'height').text = str(labeldicts[0]['img_h'])  
    ET.SubElement(sizes, 'depth').text = str(3)                    #图片的通道数：img.shape[2]
    for labeldict in labeldicts:
        objects = ET.SubElement(root, 'object')                 #创建object子节点
        ET.SubElement(objects, 'name').text = labeldict['name']        #BDD100K_10.names文件中  
                                                                       #的类别名
        ET.SubElement(objects, 'pose').text = 'Unspecified'
        ET.SubElement(objects, 'truncated').text = '0'
        ET.SubElement(objects, 'difficult').text = '0'
        bndbox = ET.SubElement(objects,'bndbox')
        ET.SubElement(bndbox, 'xmin').text = str(int(labeldict['xmin']))
        ET.SubElement(bndbox, 'ymin').text = str(int(labeldict['ymin']))
        ET.SubElement(bndbox, 'xmax').text = str(int(labeldict['xmax']))
        ET.SubElement(bndbox, 'ymax').text = str(int(labeldict['ymax']))
    tree = ET.ElementTree(root)
    tree.write(filepath, encoding='utf-8')


def save_xml(oir_name, boxes, h_w):
    qx_name = ['bianyichang', 'jiaoyichang', 'baisedian', 'shensediankuai', 'qiansekuai', 'guangquan']
    annotations_path = '/media/hjh/workdir/0_Deep_Learning/TianChi/tile_round1_train_20201231/test_set/'
    
    name = os.path.splitext(os.path.basename(oir_name))[0] + '_test'
    labeldicts = []
    for i in boxes:
        cls_id = i[0]
        bbox = i[1:]
        img_h = h_w[0]
        img_w = h_w[1]
        labeldict = {
            'name': qx_name[cls_id],
            'difficult': '0',
            'xmin': bbox[0],
            'ymin': bbox[1],
            'xmax': bbox[2],
            'ymax': bbox[3],
            'img_h': img_h,
            'img_w': img_w
        }
        labeldicts.append(labeldict)
    write_xml(name + '.jpg', annotations_path + name + '.xml', labeldicts)
    print(str( name + '.xml')+' 生成成功')
